from fastapi import FastAPI, WebSocket, Depends, WebSocketDisconnect, WebSocketException
from config.settings import get_settings, Settings
from extends.redis import async_redis_connection
from get_logger import get_websocket_logger, get_error_logger
from fastapi.responses import HTMLResponse
from contextlib import asynccontextmanager
from extends.email import thread_pool_send_email
import uuid
import redis
import orjson
import asyncio

"""
    websocket server
"""

settings: Settings = get_settings()

websocket_logger = get_websocket_logger()
error_logger = get_error_logger()

@asynccontextmanager
async def lifespan(app: FastAPI):
    try:
        yield
    finally:
        # 首先关闭所有活动的WebSocket连接
        close_tasks = []
        for websocket in active_connections.copy():
            if websocket in stop_flags:
                stop_flags[websocket].set()
            if websocket in task_running_flags:
                task_running_flags[websocket].clear()
            close_tasks.append(websocket.close())
        
        if close_tasks:
            # 等待所有WebSocket连接关闭，设置超时
            try:
                await asyncio.wait_for(asyncio.gather(*close_tasks, return_exceptions=True), timeout=3.0)
                websocket_logger.info("All websocket connections closed")
            except asyncio.TimeoutError:
                websocket_logger.warning("Some websocket connections failed to close")
        
        # 最后关闭Redis连接
        try:
            await async_redis_connection.close()
            websocket_logger.info("Redis connection closed")
        except asyncio.TimeoutError:
            websocket_logger.warning("Redis connection close timeout")
        except asyncio.CancelledError:
            websocket_logger.warning("Redis connection close cancelled")
        except Exception as e:
            error_logger.error(f"Error closing redis connection: {str(e)}")
        

app = FastAPI(lifespan=lifespan)

active_connections = set()
stop_flags = {}
task_running_flags = {}

html = """
<!DOCTYPE html>
<html>
    <head>
        <title>Chat</title>
    </head>
    <body>
        <h1>hhs data send to console</h1>
        <script>
            var ws = new WebSocket("ws://localhost:{}");
            // 接受服务端数据
            ws.onopen = function(event) {{
                ws.send("get_status");
                console.log("Sent request for data");
            }};
            ws.onmessage = function(event) {{
                console.log(JSON.parse(event.data))
            }};
            ws.onclose = function(event) {{
                ws.send("stop_get_status");
                console.log("Sent stop request");
            }};
        </script>
    </body>
</html>
""".format(settings.WEBSOCKET_PORT)

print(f"-------Test WebSocket Date: http://127.0.0.1:{settings.WEBSOCKET_PORT}/ws----------")

@app.get("/ws")
async def get():
    return HTMLResponse(html)

async def handle_send_data(websocket: WebSocket, redis_conn: redis.Redis, stop_flag: asyncio.Event, task_running: asyncio.Event):
    """处理发送给客户端的数据"""
    current_result = {}
    previous_result = {}
    while not stop_flag.is_set():
        if task_running.is_set():  
            data = await async_redis_connection.async_get_key(settings.DATA_RESULT_KEY)
            if data == "exit":
                stop_flags[websocket].set()
                task_running_flags[websocket].clear()
                await websocket.close()
                websocket_logger.info(f"[{websocket.connection_id}] closed")
                break
            if data:
                current_result = orjson.loads(data)
                if current_result != previous_result:
                    await websocket.send_json(current_result)  
                    previous_result = current_result
                    websocket_logger.info(f"[{websocket.connection_id}] send data")
            await asyncio.sleep(settings.GET_DATA_FREQUENCY)
        else:  
            await task_running.wait()


@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket, redis_conn: redis.Redis = Depends(async_redis_connection.async_get_redis_connection)):
    websocket.connection_id = str(uuid.uuid4())
    active_connections.add(websocket)
    stop_flags[websocket] = asyncio.Event()
    task_running_flags[websocket] = asyncio.Event()  

    await websocket.accept()
    send_data_task = None
    websocket_logger.info(f"[{websocket.connection_id}] connect")
    if settings.ENABLED_EMAIL:
        thread_pool_send_email(settings.TO_ADDR, f"PHD-W Client connected.", f"[{websocket.connection_id}] PHD-W client has connected to the WebSocket.")
    try:
        while True:
            data = await websocket.receive_text()
            if data == "get_status":
                if not task_running_flags[websocket].is_set():  
                    # 首次发送get_status
                    stop_flags[websocket].clear()  
                    task_running_flags[websocket].set()  
                    send_data_task = asyncio.create_task(handle_send_data(websocket, redis_conn, stop_flags[websocket], task_running_flags[websocket]))
                else:
                    # 重复发送get_status
                    data = await async_redis_connection.async_get_key(settings.DATA_RESULT_KEY)
                    if data:
                        current_result = orjson.loads(data)
                        await websocket.send_json(current_result)  
                        websocket_logger.info(f"[{websocket.connection_id}] send one data")
                    await asyncio.sleep(settings.GET_DATA_FREQUENCY)

            elif data == "stop_get_status":
                stop_flags[websocket].set()  
                task_running_flags[websocket].clear()  
                
            elif data == "close":
                stop_flags[websocket].set()
                task_running_flags[websocket].clear()
                await websocket.close()
                websocket_logger.info(f"[{websocket.connection_id}] closed")
                break

    except (WebSocketDisconnect, WebSocketException):
        websocket_logger.warning(f"[{websocket.connection_id}] disconnect")
        if settings.ENABLED_EMAIL:
            thread_pool_send_email(settings.TO_ADDR, f"PHD-W Client disconnected.", f"[{websocket.connection_id}] PHD-W client has disconnected from the WebSocket.")
    except Exception as e:
        error_logger.error(e)
    finally:
        if send_data_task and not send_data_task.done():  
            send_data_task.cancel()
        active_connections.remove(websocket)
        del stop_flags[websocket]
        del task_running_flags[websocket]  
        websocket_logger.info(f"[{websocket.connection_id}] exit")

    return None


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=settings.WEBSOCKET_HOST, port=settings.WEBSOCKET_PORT)